#include <iostream>
#include <fstream>
#include <string>
#include <chrono>
#include <thread>
#include <unistd.h>
#include <libgen.h>
#include <signal.h>

/* ROS */

#include <ros/ros.h>
#include <ros/rate.h>

#include <std_msgs/Header.h>
#include <sensor_msgs/Image.h>
#include <stereo_msgs/DisparityImage.h>
#include <sensor_msgs/Imu.h>

/* Bring in gtest */

#include <gtest/gtest.h>

std::shared_ptr<ros::NodeHandle> nh;
ros::Subscriber sub;

int frameCount;
int droppedFramesCount;
int lastFrameSeq;

void commonCallback(const std_msgs::Header hdr)
{

    if (frameCount && lastFrameSeq)
    {
        droppedFramesCount += (hdr.seq - lastFrameSeq) - 1;
    }

    lastFrameSeq = hdr.seq;

    frameCount++;
}

void imuCallback(const sensor_msgs::ImuConstPtr msg)
{
    commonCallback(msg->header);
}

void disparityCallback(const stereo_msgs::DisparityImagePtr msg)
{
    commonCallback(msg->header);
}

void imageCallback(const sensor_msgs::ImageConstPtr msg)
{
    commonCallback(msg->header);
}

void TestImageTopic(const float expectedFrameRate)
{
        ros::Time beginTime = ros::Time::now();

        frameCount = 0;

        for (int i=0; i<100 && !frameCount; i++)
        {
            ros::Duration(0.1).sleep();
        }

        ASSERT_GT(frameCount, 0);

        frameCount = 0;
        droppedFramesCount = 0;

        ros::Time firstFrameTime = ros::Time::now();

        ros::Duration d = beginTime - firstFrameTime;

        ASSERT_LT(d.toNSec(), 1000000ll);

        ros::Duration waitTime(5.0f);

        waitTime.sleep();

        ros::Time lastFrameTime = ros::Time::now();

        sub.shutdown();

        //ROS_INFO_STREAM("frame count = " << frameCount);
        EXPECT_GT(frameCount, waitTime.toSec() * expectedFrameRate * 0.8f);

        float fr = 1.0e9 * frameCount / (lastFrameTime - firstFrameTime).toNSec();
        //ROS_INFO_STREAM("Frame rate = " << fr);
        EXPECT_GT(fr, expectedFrameRate * 0.95f);
        EXPECT_LT(fr, expectedFrameRate * 1.05f);

        //ROS_INFO_STREAM("drop count = " << droppedFramesCount);
        EXPECT_EQ(droppedFramesCount, 0);
}

TEST(TestSuite, failTest)
{
    EXPECT_GT(0, 1);
}

TEST(TestSuite, webcamFrames)
{
    try
    {
        sub.shutdown();
        sub = nh->subscribe("/sensor_msgs/Image/Webcam" , 5, imageCallback);

        EXPECT_NE(sub, nullptr);

        TestImageTopic(30.0f);
    }
    catch(...)
    {
        ADD_FAILURE() << "Uncaught exception";
    }
}

TEST(TestSuite, videoFrames)
{ 
    try
    {   
        sub.shutdown();
        sub = nh->subscribe("/sensor_msgs/Image/Video/left/image" , 5, imageCallback);

        EXPECT_NE(sub, nullptr);

        TestImageTopic(30.0f);
    }
    catch(...)
    {   
        ADD_FAILURE() << "Uncaught exception";
    }
}

TEST(TestSuite, depthFrames)
{
    try
    {
        sub.shutdown();
        sub = nh->subscribe("/sensor_msgs/Image/Depth", 5, imageCallback);

        EXPECT_NE(sub, nullptr);

        TestImageTopic(30.0f);
    }
    catch(...)
    {
        ADD_FAILURE() << "Uncaught exception";
    }
}

TEST(TestSuite, disparityFrames)
{
    try
    {
        sub.shutdown();
        sub = nh->subscribe("/sensor_msgs/Image/Video/disparity" , 5, disparityCallback);

        EXPECT_NE(sub, nullptr);

        TestImageTopic(30.0f);
    }
    catch(...)
    {
        ADD_FAILURE() << "Uncaught exception";
    }
}

TEST(TestSuite, imuFrames)
{
    try
    {
        sub.shutdown();
        sub = nh->subscribe("/imu/sensor_msgs/Imu" , 5, imuCallback);

        EXPECT_NE(sub, nullptr);

        TestImageTopic(225.0f);
    }
    catch(...)
    {
        ADD_FAILURE() << "Uncaught exception";
    }
}

TEST(TestSuite, webcamPublishingStartStop)
{   
    try
    {
        int frameWaitLoopCount;

        sub.shutdown();

        for (int cnt=0;cnt<10;cnt++)
        {
            frameCount = 0;

            sub = nh->subscribe("/sensor_msgs/Image/Webcam" , 5, imageCallback);

            EXPECT_NE(sub, nullptr);

            for (frameWaitLoopCount=100; frameWaitLoopCount && !frameCount; frameWaitLoopCount--)
            {
                ros::Duration(0.1).sleep();
            }

            sub.shutdown();

            ASSERT_NE(frameWaitLoopCount, 0);
            ASSERT_GT(frameCount, 0);

            frameCount = 0;

            for (frameWaitLoopCount=10; frameWaitLoopCount && !frameCount; frameWaitLoopCount--)
            {
                ros::Duration(0.1).sleep();
            }

            ASSERT_EQ(frameWaitLoopCount, 0);
            ASSERT_EQ(frameCount, 0);
        }
    }
    catch(...)
    {   
        ADD_FAILURE() << "Uncaught exception";
    }
}

TEST(TestSuite, videoPublishingStartStop)
{
    try
    {
        int i;

        sub.shutdown();

        for (int cnt=0;cnt<10;cnt++)
        {
            frameCount = 0;

            sub = nh->subscribe("/sensor_msgs/Image/Video/left/image", 5, imageCallback);

            EXPECT_NE(sub, nullptr);

            for (i=100; i && !frameCount; i--)

            for (i=100; i && !frameCount; i--)
            {
                ros::Duration(0.1).sleep();
            }

            sub.shutdown();

            ASSERT_NE(i, 0);
            ASSERT_GT(frameCount, 0);

            frameCount = 0;

            for (i=10; i && !frameCount; i--)
            {
                ros::Duration(0.1).sleep();
            }

            ASSERT_EQ(i, 0);
            ASSERT_EQ(frameCount, 0);
        }
    }
    catch(...)
    {
        ADD_FAILURE() << "Uncaught exception";
    }
}

TEST(TestSuite, depthPublishingStartStop)
{
    try
    {
        int i;

        sub.shutdown();

        for (int cnt=0;cnt<10;cnt++)
        {
            frameCount = 0;

            sub = nh->subscribe("/sensor_msgs/Image/Depth", 5, imageCallback);

            EXPECT_NE(sub, nullptr);

            for (i=100; i && !frameCount; i--)
            {   
                ros::Duration(0.1).sleep();
            }

            sub.shutdown();

            ASSERT_NE(i, 0);
            ASSERT_GT(frameCount, 0);

            frameCount = 0;

            for (i=10; i && !frameCount; i--)
            {   
                ros::Duration(0.1).sleep();
            }

            ASSERT_EQ(i, 0);
            ASSERT_EQ(frameCount, 0);
        }
    }
    catch(...)
    {
        ADD_FAILURE() << "Uncaught exception";
    }
}

TEST(TestSuite, disparityPublishingStartStop)
{
    try
    {
        int i;

        sub.shutdown();

        for (int cnt=0;cnt<10;cnt++)
        {
            frameCount = 0;

            sub = nh->subscribe("/sensor_msgs/Image/Video/disparity", 5, disparityCallback);

            EXPECT_NE(sub, nullptr);

            for (i=100; i && !frameCount; i--)
            {
                ros::Duration(0.1).sleep();
            }

            sub.shutdown();

            ASSERT_NE(i, 0);
            ASSERT_GT(frameCount, 0);

            frameCount = 0;

            for (i=10; i && !frameCount; i--)
            {
                ros::Duration(0.1).sleep();
            }

            ASSERT_EQ(i, 0);
            ASSERT_EQ(frameCount, 0);
        }
    }
    catch(...)
    {
        ADD_FAILURE() << "Uncaught exception";
    }
}

TEST(TestSuite, imuPublishingStartStop)
{
    try
    {
        int i;

        sub.shutdown();

        for (int cnt=0;cnt<10;cnt++)
        {
            frameCount = 0;

            sub = nh->subscribe("/imu/sensor_msgs/Imu", 5, imuCallback);
    
            EXPECT_NE(sub, nullptr);

            for (i=100; i && !frameCount; i--)
            {
                ros::Duration(0.1).sleep();
            }

            sub.shutdown();

            ASSERT_NE(i, 0);
            ASSERT_GT(frameCount, 0);

            frameCount = 0;

            for (i=10; i && !frameCount; i--)
            {
                ros::Duration(0.1).sleep();
            }

            ASSERT_EQ(i, 0);
            ASSERT_EQ(frameCount, 0);
        }
    }
    catch(...)
    {
        ADD_FAILURE() << "Uncaught exception";
    }
}

// Run all the tests that were declared with TEST()
int main(int argc, char **argv)
{
    ros::init(argc, argv, "tester");
    nh.reset(new ros::NodeHandle);

    testing::InitGoogleTest(&argc, argv);

    std::thread t([]
    {
        while(ros::ok())
        {
            ros::spin();
        }
    });

    auto res = RUN_ALL_TESTS();

    ros::shutdown();

    t.join();

    return res;
}
